#include "NaluParse.h"

#define H264_TYPE(v) ((uint8_t)(v) & 0x1F)
#define H265_TYPE(v) (((uint8_t)(v) >> 1) & 0x3f)

static int32_t _startLen(const uint8_t *in)
{
    int32_t start_ = 0;
    if(in[0] == 0x0){
        if(in[2] == 0x01) start_ = 3;
        else if(in[2] == 0x00 && in[3] == 0x01) start_ = 4;
        else start_ = -1;
    }
    return start_;
}

char keyframe(NALU_t *nalu)
{
    if(nalu == NULL)return NaluFalse;
    int32_t nalu_type = nalu->nalType;
    if(nalu->type == 0){
        return (nalu_type == NAL_UNIT_TYPE_CODED_SLICE_IDR && decodeAble(nalu));
    }else if(nalu->type == 1){
        return ((nalu_type >= NAL_UNIT_CODED_SLICE_BLA_W_LP && 
                 nalu_type <= NAL_UNIT_RESERVED_IRAP_VCL23) && 
                 decodeAble(nalu));
    }
    return NaluFalse;
}

char _keyframe(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == 0){
        int8_t nalu_type = H264_TYPE(in[start_]);
        return (nalu_type == 
            NAL_UNIT_TYPE_CODED_SLICE_IDR) && _decodeAble(type, in, size, startLen_);
    }else if(type == 1){
        int8_t nalu_type = H265_TYPE(in[start_]);
        return ((nalu_type >= NAL_UNIT_CODED_SLICE_BLA_W_LP && 
                 nalu_type <= NAL_UNIT_RESERVED_IRAP_VCL23) && 
                 _decodeAble(type, in, size, startLen_));
    }
    return NaluFalse;
}

char configFrame(NALU_t *nalu)
{
    if(nalu == NULL)return NaluFalse;
    int32_t nalu_type = nalu->nalType;
    if(nalu->type == 0){
        if(nalu_type == NAL_UNIT_TYPE_SPS || 
           nalu_type == NAL_UNIT_TYPE_PPS) return NaluTrue;
    }else if(nalu->type == 1){
        if(nalu_type == NAL_UNIT_VPS || 
           nalu_type == NAL_UNIT_SPS || 
           nalu_type == NAL_UNIT_PPS) return NaluTrue;
    }
    return NaluFalse;
}

char _configFrame(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == 0){
        int8_t nalu_type = H264_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_TYPE_SPS || 
           nalu_type == NAL_UNIT_TYPE_PPS) return NaluTrue;
    }else if(type == 1){
        int8_t nalu_type = H265_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_VPS || 
           nalu_type == NAL_UNIT_SPS || 
           nalu_type == NAL_UNIT_PPS) return NaluTrue;
    }
    return NaluFalse;
}

char dropAble(NALU_t *nalu)
{
    if(nalu == NULL)return NaluFalse;
    int32_t nalu_type = nalu->nalType;
    if(nalu->type == 0){
        if(nalu_type == NAL_UNIT_TYPE_SEI || 
           nalu_type == NAL_UNIT_TYPE_AUD) return NaluTrue;
    }else{
        if(nalu_type == NAL_UNIT_PREFIX_SEI || 
           nalu_type == NAL_UNIT_SUFFIX_SEI || 
           nalu_type == NAL_UNIT_AUD) return NaluTrue;
    }
    return NaluFalse;
}

char _dropAble(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == 0){
        int8_t nalu_type = H264_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_TYPE_SEI || 
           nalu_type == NAL_UNIT_TYPE_AUD) return NaluTrue;
    }else if(type == 1){
        int8_t nalu_type = H265_TYPE(in[start_]);
        if(nalu_type == NAL_UNIT_PREFIX_SEI || 
           nalu_type == NAL_UNIT_SUFFIX_SEI || 
           nalu_type == NAL_UNIT_AUD) return NaluTrue;
    }
    return NaluFalse;
}

char decodeAble(NALU_t *nalu)
{
    if(nalu == NULL)return NaluFalse;
    int32_t nalu_type = nalu->nalType;
    if(nalu->type == 0){
        return (nalu_type >= NAL_UNIT_TYPE_CODED_SLICE_NON_IDR && 
                nalu_type <= NAL_UNIT_TYPE_CODED_SLICE_IDR && 
                (nalu->first_mb_in_slice & 0x80));
    }else if(nalu->type == 1){
        return (nalu_type >= NAL_UNIT_CODED_SLICE_TRAIL_N && 
                nalu_type <= NAL_UNIT_RESERVED_IRAP_VCL23 && 
                (nalu->first_mb_in_slice & 0x80));
    }
    return NaluFalse;
}

char _decodeAble(int type, const uint8_t *in, size_t size, int32_t *startLen_)
{
    *startLen_ = 0;
    if(in == NULL || size < 5)return NaluFalse;
    int32_t start_ = _startLen(in);
    if(start_ == -1)return NaluFalse;
    *startLen_ = start_;
    if(type == 0){
        int8_t nalu_type = H264_TYPE(in[start_]);
        return (nalu_type >= NAL_UNIT_TYPE_CODED_SLICE_NON_IDR && 
                nalu_type <= NAL_UNIT_TYPE_CODED_SLICE_IDR && 
                (in[start_+1] & 0x80));
    }else if(type == 1){
        int8_t nalu_type = H265_TYPE(in[start_]);
        return (nalu_type >= NAL_UNIT_CODED_SLICE_TRAIL_N && 
                nalu_type <= NAL_UNIT_RESERVED_IRAP_VCL23 && 
                (in[start_+2] & 0x80));
    }
    return NaluFalse;
}